
import re
import requests
import json
import base64



def test_api_direct_V2(image_path, task):    
    base_url = "https://envision-research.hkust-gz.edu.cn/vlm_cloth_api_v0.2"
    api_url = f"{base_url}/api/predict"
        
    with open(image_path, "rb") as f:
        image_data = base64.b64encode(f.read()).decode()
    payload = {
        "data": [
            f"data:image/png;base64,{image_data}",
            task
        ],
        "fn_index": 0
    }
    
    response = requests.post(api_url, json=payload)
    raw_response = response.json()['data'][0]
    print(raw_response)
    try:
        if isinstance(raw_response, str):
            action_data = json.loads(raw_response)
        else:
            action_data = raw_response
    except json.JSONDecodeError:
        raw_text = str(raw_response)
        
        def extract_arm_section(text, arm_name):
            pattern = f'{arm_name}:' + r'\{'
            start_match = re.search(pattern, text)
            if not start_match:
                return ""
            
            start_pos = start_match.end() - 1
            brace_count = 0
            pos = start_pos
            
            while pos < len(text):
                if text[pos] == '{':
                    brace_count += 1
                elif text[pos] == '}':
                    brace_count -= 1
                    if brace_count == 0:
                        return text[start_pos + 1:pos]
                pos += 1
            
            return ""
        
        left_arm_text = extract_arm_section(raw_text, 'left_arm')
        right_arm_text = extract_arm_section(raw_text, 'right_arm')
        
        action_data = {
            "left_arm": left_arm_text,
            "right_arm": right_arm_text
        }
    
    PICK_RE  = re.compile(r"pick:\s*\{[^}]*?point_2d:\s*\[\s*(\d+)\s*,\s*(\d+)\s*\]",  re.I | re.S)
    PLACE_RE = re.compile(r"place:\s*\{[^}]*?point_2d:\s*\[\s*(\d+)\s*,\s*(\d+)\s*\]", re.I | re.S)
    
    def extract_coordinates(arm_data):
        text_data = str(arm_data)
        
        pick_match = PICK_RE.search(text_data)
        if pick_match:
            pick_coords = [int(pick_match.group(1)), int(pick_match.group(2))]
        else:
            pick_coords = [0, 0]
        
        # Check if there's "no need place" in the text
        if "no need place" in text_data.lower():
            place_coords = [0, 0]
        else:
            place_match = PLACE_RE.search(text_data)
            if place_match:
                place_coords = [int(place_match.group(1)), int(place_match.group(2))]
            else:
                place_coords = [0, 0]
        
        return {
            "pick": pick_coords,
            "place": place_coords
        }
    
    final_result = {
        "action": {
            "left_arm": extract_coordinates(action_data.get("left_arm", action_data)),
            "right_arm": extract_coordinates(action_data.get("right_arm", action_data))
        },
        "task": task
    }
    
    print(f"Parsed Result: {json.dumps(final_result, indent=2, ensure_ascii=False)}")
    return final_result